import os

import pandas as pd
from pandas import DataFrame

from app_config import get_engine_ts, get_pro

"""
csi all exchange_amount > 300 yi
exclude hs_300 
exclude zz_500
exclude sz_50
"""
engine = get_engine_ts()


def calc(date_sz50: str, strStartDate: str, strEndDate: str):
    df_csi_all = get_pro().index_weight(index_code='000985.CSI', start_date=date_sz50, end_date=date_sz50)
    list_csi_all_code = df_csi_all['con_code'].tolist()

    query = f"""
       SELECT ts_code,trade_date, amount FROM `daily`
       WHERE ts_code IN ({','.join(f"'{code}'" for code in list_csi_all_code)})
       AND trade_date >= '{strStartDate}'
       AND trade_date <= '{strEndDate}'
       """

    # 执行查询并将结果转换为DataFrame
    df_daily_csi_all = pd.read_sql_query(query, engine)
    df_daily_csi_all['amount'] = df_daily_csi_all['amount'] / 1000
    grouped_df = df_daily_csi_all.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', 'avg_amount']
    merge = pd.merge(df_csi_all, grouped_df, left_on='con_code', right_on='ts_code', how='left')
    # merge.to_excel('./csi_all.xlsx', index=False)

    # 构建查询语句
    query = f"""
          SELECT ts_code,trade_date, total_mv FROM `daily_basic`
          WHERE ts_code IN ({','.join(f"'{code}'" for code in list_csi_all_code)})
          AND trade_date >= '{strStartDate}'
          AND trade_date <= '{strEndDate}'
          """
    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['total_mv'] = result_df['total_mv'] / 10000
    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', 'avg_total_mv']
    pd_merge = pd.merge(merge, grouped_df_total_mv, left_on='con_code', right_on='ts_code', how='left')
    pd_merge.to_excel('./csi_all1.xlsx', index=False)


def te():
    return 1 == 1


def calc1(date_sz50: str):
    df_csi_all: DataFrame = pd.read_excel('./csi_all1.xlsx')
    df_stock_basic = df_csi_all[((df_csi_all['avg_total_mv'] > 200.0) & (True))]
    # df_stock_basic.to_excel('./stock_basic.xlsx', index=False)
    index_basics: list[str] = ['000016.SH', '000300.SH', '399905.SZ', '000510.CSI', '399006.SZ', '399673.SZ', '000688.SH', '000852.SH']
    dict_index_name: dict = {}

    for index_basic in index_basics:
        query = f"""
        SELECT distinct ts_code,name FROM `index_basic` WHERE ts_code = '{index_basic}'
         """
        # 执行查询并将结果转换为DataFrame
        result_df = pd.read_sql_query(query, engine)
        index_name: str = result_df['name'].tolist()[0]

        df_index_basic = get_pro().index_weight(index_code=index_basic, start_date=date_sz50, end_date=date_sz50)
        con_codes = df_index_basic['con_code'].tolist()
        for con_code in con_codes:

            if dict_index_name.__contains__(con_code):
                dict_index_name[con_code] = dict_index_name[con_code] + ', ' + index_name
            else:
                dict_index_name[con_code] = index_name

    query = f"""
           SELECT distinct ts_code,name FROM `stock_basic`
            """
    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    df_csi_all['value'] = df_csi_all['con_code'].map(dict_index_name)
    df_csi_all = df_csi_all.merge(result_df, left_on='con_code', right_on='ts_code', how='left')
    df_csi_all.to_excel('./csi_all2.xlsx')


if __name__ == '__main__':
    calc(date_sz50='20250430', strStartDate='20240501', strEndDate="20250430")
    calc1(date_sz50="20250430")
